use syn::{self, visit_mut::VisitMut, Item, ItemFn, ItemStruct, ItemEnum, ItemStatic, ItemConst, Fields, File, Visibility};
use quote::ToTokens;
use walkdir::WalkDir;
use std::fs;
use std::io::Write;
use std::env;
use std::collections::HashSet;

struct PubAdder {
    file_name: String,
}

impl VisitMut for PubAdder {
    fn visit_item_fn_mut(&mut self, item: &mut ItemFn) {
        if self.file_name == "common_interface_h.rs" {
            // Remove all functions in common_interface_h.rs
            *item = ItemFn {
                attrs: vec![],
                vis: Visibility::Inherited,
                sig: syn::parse_quote!(fn __removed_function()),
                block: Box::new(syn::parse_quote!({})),
            };
        } else {
            if item.block.stmts.is_empty() {
                // Remove functions with empty bodies
                *item = ItemFn {
                    attrs: vec![],
                    vis: Visibility::Inherited,
                    sig: syn::parse_quote!(fn __removed_function()),
                    block: Box::new(syn::parse_quote!({})),
                }
            } else {
                item.vis = Visibility::Public(syn::token::Pub::default());
            }
        }
        syn::visit_mut::visit_item_fn_mut(self, item);
    }

    fn visit_item_struct_mut(&mut self, item: &mut ItemStruct) {
        item.vis = Visibility::Public(syn::token::Pub::default());
        if let Fields::Named(ref mut fields) = item.fields {
            for field in &mut fields.named {
                field.vis = Visibility::Public(syn::token::Pub::default());
            }
        }
        syn::visit_mut::visit_item_struct_mut(self, item);
    }

    fn visit_item_enum_mut(&mut self, item: &mut ItemEnum) {
        item.vis = Visibility::Public(syn::token::Pub::default());
        syn::visit_mut::visit_item_enum_mut(self, item);
    }

    fn visit_item_static_mut(&mut self, item: &mut ItemStatic) {
        item.vis = Visibility::Public(syn::token::Pub::default());
        syn::visit_mut::visit_item_static_mut(self, item);
    }

    fn visit_item_const_mut(&mut self, item: &mut ItemConst) {
        item.vis = Visibility::Public(syn::token::Pub::default());
        syn::visit_mut::visit_item_const_mut(self, item);
    }
}

fn filter_removed_functions(file: &mut File) {
    file.items.retain(|item| {
        if let Item::Fn(item_fn) = item {
            if item_fn.sig.ident == "__removed_function" {
                return false;
            }
        }
        true
    })
}

fn extract_globals(file: &File) -> HashSet<String> {
    let mut globals = HashSet::new();
    for item in &file.items {
        match item {
            Item::Static(item_static) => {
                globals.insert(format!("{}: {}", item_static.ident, item_static.ty.to_token_stream()));
            },
            Item::Const(item_const) => {
                globals.insert(format!("{}: {}", item_const.ident, item_const.ty.to_token_stream()));
            }
            _ => {}
        }
    }
    globals
}

fn remove_duplicate_globals(file: &mut File, common_globals: &HashSet<String>) {
    file.items.retain(|item| {
        match item {
            Item::Static(item_static) => {
                let global_decl = format!("{}: {}", item_static.ident, item_static.ty.to_token_stream());
                !common_globals.contains(&global_decl)
            },
            Item::Const(item_const) => {
                let global_decl = format!("{}: {}", item_const.ident, item_const.ty.to_token_stream());
                !common_globals.contains(&global_decl)
            },
            _ => true
        }
    })
}

fn process_file(file_path: &str, file_name: &str, all_files: &[String], common_globals: &HashSet<String>) {
    let content = fs::read_to_string(file_path).expect("Unable to read file");
    println!("Parsing {}", file_path);
    let mut syntax_tree: File = syn::parse_file(&content).expect("Unable to parse file");

    let mut pub_adder = PubAdder {
        file_name: file_name.to_string(),
    };
    pub_adder.visit_file_mut(&mut syntax_tree);

    if file_name != "common_interface_h.rs" {
        remove_duplicate_globals(&mut syntax_tree, common_globals);
    }

    filter_removed_functions(&mut syntax_tree);

    let mut new_content = syntax_tree.into_token_stream().to_string();

    // syn and quote can produce code that doesn't have the same formatting as the original
    // code. Using a formatter like rustfmt can help ensure consistent formatting.
    // Assuming `rustfmt` is installed in the environment:
    if let Ok(mut output) = std::process::Command::new("rustfmt").arg("--emit").arg("stdout").stdin(std::process::Stdio::piped()).stdout(std::process::Stdio::piped()).spawn() {
        if let Some(stdin) = output.stdin.as_mut() {
            write!(stdin, "{}", new_content).expect("Failed to write to rustfmt stdin");
        }
        if let Ok(output) = output.wait_with_output() {
            if output.status.success() {
                new_content = String::from_utf8_lossy(&output.stdout).into_owned();
            }
        }
    }

    if file_name == "main.rs" {
        let mut main_mods = String::new();
        for file in all_files {
            if file != "main.rs" {
                let mod_name = file.strip_suffix(".rs").unwrap();
                main_mods.push_str(&format!("mod {};\n", mod_name));
            }
        }
        for file in all_files {
            if file != "main.rs" {
                let mod_name = file.strip_suffix(".rs").unwrap();
                main_mods.push_str(&format!("use crate::{}::*;\n", mod_name));
            }
        }
        new_content = format!("{}\n{}", main_mods, new_content);
    } else if file_name != "common_interface_h.rs" {
        new_content = format!("use crate::*;\n{}", new_content);
    }

    fs::write(file_path, new_content).expect("Unable to write file");
}

fn main() {
    let args: Vec<String> = env::args().collect();

    // Cargo project name

    if args.len() != 2 {
        println!("Usage: {} <name_to_cargo_project>", args[0]);
    }

    let project_name = &args[1];
    //let src_path = format!("{}/src", project_name);
    let src_path = project_name.clone();

    let mut all_files = vec![];
    let mut common_globals = HashSet::new();

    for entry in WalkDir::new(src_path.clone()) {
        let entry = entry.expect("Error reading direcctory entry");
        if entry.path().extension().map(|s| s == "rs").unwrap_or(false) {
            if let Some(file_name) = entry.path().file_name().and_then(|s| s.to_str()) {
                all_files.push(file_name.to_string());
            }
        }
    }

    for entry in WalkDir::new(src_path.clone()) {
        let entry = entry.expect("Error reading directory entry");
        if entry.path().extension().map(|s| s == "rs").unwrap_or(false) {
            let file_path = entry.path().to_str().expect("Error converting path to string");
            let file_name = entry.path().file_name().and_then(|s| s.to_str()).expect("Error getting file name");
            if file_name == "common_interface_h.rs" {
                let content = fs::read_to_string(file_path).expect("Unable to read file");
                let syntax_tree: File = syn::parse_file(&content).expect("Unable to parse file");
                common_globals = extract_globals(&syntax_tree);
            }
        }
    }

    for entry in WalkDir::new(src_path.clone()) {
        let entry = entry.expect("Error reading directory entry");
        if entry.path().extension().map(|s| s == "rs").unwrap_or(false) {
            let file_path = entry.path().to_str().expect("Error converting path to string");
            let file_name = entry.path().file_name().and_then(|s| s.to_str()).expect("Error getting file name");
            process_file(file_path, file_name, &all_files, &common_globals);
        }
    }
}

